import type { TreeDataNode, TreeProps } from 'antd';
import { Tree } from 'antd';
import type { Dispatch, Key, ReactElement, Ref, RefAttributes, SetStateAction } from 'react';

type RcTree = typeof Tree extends (props: RefAttributes<infer P>) => ReactElement ? P : never;

type TreeComponentProps = {
  ref: Ref<RcTree>;
  treeData: TreeDataNode[];
  expandedKeys: Key[];
  setExpandedKeys: Dispatch<SetStateAction<Key[]>>;
  selectedKeys?: Key[];
  setSelectedKeys?: Dispatch<SetStateAction<Key[]>>;
  onSelect?: TreeProps['onSelect'];
};

type LoopCallback = (node: TreeDataNode, i: number, data: TreeDataNode[]) => void;

export const TreeComponent = (props: TreeComponentProps) => {
  const treeData = props.treeData;
  const { ref, expandedKeys, setExpandedKeys, selectedKeys, onSelect } = props;

  const onDragEnter: TreeProps['onDragEnter'] = (info) => {
    console.log(info);
    // expandedKeys, set it when controlled is needed
    // setExpandedKeys(info.expandedKeys)
  };

  const loop = (data: TreeDataNode[], key: Key, callback: LoopCallback): TreeDataNode | undefined => {
    for (let i = 0; i < data.length; i++) {
      const item = data[i];
      if (item.key === key) {
        callback(item, i, data);
        return item;
      }
      if (item.children) {
        return loop(item.children, key, callback);
      }
    }
  };

  const onDrop: TreeProps['onDrop'] = (info) => {
    console.log(info);
    const dropKey = info.node.key;
    const dragKey = info.dragNode.key;
    const dropPos = info.node.pos.split('-');
    const dropPosition = info.dropPosition - Number(dropPos[dropPos.length - 1]); // the drop position relative to the drop node, inside 0, top -1, bottom 1

    const data = [...treeData];

    // Find dragObject
    const dragObj = loop(data, dragKey, (_item, index, arr) => {
      arr.splice(index, 1);
    }) as TreeDataNode;

    if (!info.dropToGap) {
      // Drop on the content
      loop(data, dropKey, (_item) => {
        _item.children = _item.children || [];
        // where to insert. New _item was inserted to the start of the array in this example, but can be anywhere
        _item.children.unshift(dragObj);
      });
    } else {
      let ar: TreeDataNode[] = [];
      let i = -1;
      loop(data, dropKey, (_item, index, arr) => {
        ar = arr;
        i = index;
      });
      if (dropPosition === -1) {
        // Drop on the top of the drop node
        ar.splice(i, 0, dragObj);
      } else {
        // Drop on the bottom of the drop node
        ar.splice(i + 1, 0, dragObj);
      }
    }
  };

  const onExpand: TreeProps['onExpand'] = (keys, info) => {
    if (!info.expanded) {
      setExpandedKeys(expandedKeys.filter((key) => key !== info.node.key));
    } else {
      setExpandedKeys([...keys, info.node.key]);
    }
  };
  // const onSelect: TreeProps['onSelect'] = (keys, info) => {
  //   console.log(keys, info);
  // };

  return (
    <div className="w-full h-[40%] bg-cyan-300 rounded-sm">
      <Tree
        ref={ref}
        className="draggable-tree"
        // defaultExpandedKeys={expandedKeys}
        draggable={{
          icon: false,
        }}
        blockNode
        onExpand={onExpand}
        expandedKeys={expandedKeys}
        selectedKeys={selectedKeys}
        onDragEnter={onDragEnter}
        onDrop={onDrop}
        onSelect={onSelect}
        treeData={treeData}
        height={200}
      />
    </div>
  );
};
